import os

import torch

from cnn_mnist.utils import image_util


def train(model, device, train_loader, optimizer, epoch, criterion):
    # 开启训练模型: 启动 Dropout 和 BatchNorm的训练行为
    model.train()

    train_loss = 0

    correct = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        # 模型、输入数据、label需要 GPU化，交给GPU计算
        data, target = data.to(device), target.to(device)

        os.makedirs(f'lenet_train/images/', exist_ok=True)

        if data.dim() == 4:
            for i, c_h_w in enumerate(data):
                image_path = f"lenet_train/images/{batch_idx}_{i}.jpg"
                image_util.save_image_tensor(c_h_w, image_path)

        # 梯度归0
        optimizer.zero_grad()
        output = model(data)

        # 使用损失函数 计算损失
        loss = criterion(output, target)

        # 方向传播，得到偏导
        loss.backward()

        # 优化，根据微分 偏导得到的 梯度，更新模型的参数
        optimizer.step()

        train_loss += loss.item()
        pred = output.argmax(dim=1, keepdim=True)

        target_view = target.view_as(pred)
        pred_target_eq = pred.eq(target_view)

        correct += pred_target_eq.sum().item()

        if batch_idx % 10 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

    train_loss /= len(train_loader.dataset)
    accuracy = 100. * correct / len(train_loader.dataset)
    print(f'\nTrain set: Average loss: {train_loss:.4f}, Accuracy: {correct}/{len(train_loader.dataset)} '
          f'({accuracy:.2f}%)\n')

    return train_loss, accuracy


def test(model, device, test_loader, criterion):
    # 禁用 Dropout 和 固定 BatchNorm
    model.eval()

    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()

            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} '
          f'({accuracy:.2f}%)\n')
    return test_loss, accuracy


# ----------------------
# 4. 预测函数
# ----------------------
def predict_digit(model, image_tensor, device='cuda'):
    """使用模型预测数字"""
    if image_tensor is None:
        return None, None

    # 将输入移到设备上
    # image_tensor = image_tensor.to(device)

    # 禁用梯度计算
    with torch.no_grad():
        # 前向传播
        output = model(image_tensor)

        # 计算概率分布
        probabilities = torch.nn.functional.softmax(output, dim=1)

        # 获取预测结果和置信度
        confidence, predicted_class = torch.max(probabilities, 1)

    return predicted_class.item(), confidence.item()
